package cn.yx.common.webflux.util;

import cn.hutool.core.collection.CollectionUtil;
import cn.yx.common.core.util.json.jackson.JsonUtil;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;

/**
 * <p>WebFlux工具类</p>
 *
 * @author Wgssmart
 */
public final class WebFluxUtil {

    private WebFluxUtil() {
    }

    /**
     * 添加请求头到ServerHttpRequest
     *
     * @param exchange
     * @param headerName
     * @param headerValue
     * @return
     */
    public static ServerHttpRequest addHeader(ServerWebExchange exchange, String headerName, String headerValue) {
        return addHeader(exchange.getRequest(), headerName, headerValue);
    }

    /**
     * 添加请求头到ServerHttpRequest。注意request.getHeaders获取的请求头为只读的
     *
     * @param request
     * @param headerName
     * @param headerValue
     * @return
     */
    public static ServerHttpRequest addHeader(ServerHttpRequest request, String headerName, String headerValue) {
        ServerHttpRequest.Builder builder = request.mutate();
        builder.header(headerName, headerValue);
        return builder.build();
    }

    /**
     * 添加请求头到ServerHttpRequest
     *
     * @param exchange
     * @param headers
     * @return
     */
    public static ServerHttpRequest addHeaders(ServerWebExchange exchange, Map<String, String> headers) {
        return addHeaders(exchange.getRequest(), headers);
    }

    /**
     * 添加请求头到ServerHttpRequest
     *
     * @param request
     * @param headers
     * @return
     */
    public static ServerHttpRequest addHeaders(ServerHttpRequest request, Map<String, String> headers) {
        ServerHttpRequest.Builder builder = request.mutate();
        for (Map.Entry<String, String> headerEntry : headers.entrySet()) {
            builder.header(headerEntry.getKey(), headerEntry.getValue());
        }
        return builder.build();
    }

    /**
     * 从ServerHttpRequest获取请求头
     *
     * @param request
     * @param headerName
     * @return
     */
    public static String getHeader(ServerHttpRequest request, String headerName) {
        List<String> headerValues = request.getHeaders().get(headerName);
        return CollectionUtil.isNotEmpty(headerValues) ? headerValues.get(0) : null;
    }

    /**
     * 从ServerHttpRequest中获取请求数据
     *
     * @param request
     * @param tClass
     * @param <T>
     * @return
     */
    public static <T> T getRequestData(ServerHttpRequest request, Class<T> tClass) {
        return Mono.from(request.getBody()).ofType(tClass).block();
    }

    /**
     * 构建{@code Mono<DataBuffer>}对象，可用于ServerHttpResponse输出响应数据
     *
     * @param response
     * @param data
     * @param charset
     * @param <T>
     * @return
     */
    public static <T> Mono<DataBuffer> monoJson(ServerHttpResponse response, T data, Charset charset) {
        DataBuffer dataBuffer = response.bufferFactory().wrap(JsonUtil.toJsonStr(data).getBytes(charset));
        return Mono.just(dataBuffer);
    }

    /**
     * 构建{@code Mono<DataBuffer>}对象，可用于ServerHttpResponse输出响应数据，字符编码默认为UTF-8
     *
     * @param response
     * @param data
     * @param <T>
     * @return
     */
    public static <T> Mono<DataBuffer> monoJson(ServerHttpResponse response, T data) {
        return monoJson(response, data, StandardCharsets.UTF_8);
    }

    /**
     * ServerHttpResponse输出响应数据
     *
     * @param response
     * @param data
     * @param charset
     * @param <T>
     * @return
     */
    public static <T> Mono<Void> write(ServerHttpResponse response, T data, Charset charset) {
        response.setStatusCode(HttpStatus.OK);
        response.getHeaders().setContentType(MediaType.APPLICATION_JSON);
        return response.writeWith(monoJson(response, data, charset));
    }

    /**
     * ServerHttpResponse输出响应数据，字符编码默认为UTF-8
     *
     * @param response
     * @param data
     * @param <T>
     * @return
     */
    public static <T> Mono<Void> write(ServerHttpResponse response, T data) {
        response.setStatusCode(HttpStatus.OK);
        response.getHeaders().setContentType(MediaType.APPLICATION_JSON);
        return response.writeWith(monoJson(response, data, StandardCharsets.UTF_8));
    }

}
